{-# LANGUAGE DoRec #-}

module HaskellParser 
(
  parseHaskellFile
)
where

import Prelude ()
import StdImports
import qualified Data.Map as Map

import Zeno

import qualified GHC.Paths as Paths
import qualified CoreSyn as Hs
import qualified Var as Hs
import qualified Type as Hs
import qualified HscTypes as Hs
import qualified DataCon as Hs
import qualified TyCon as Hs
import qualified GHC as Hs
import qualified DynFlags as Hs

import Outputable ( Outputable, ppr, defaultUserStyle )


data HsEnv = HsEnv  { envVars :: Map Hs.Var ZVar,
                      envTypes :: Map Hs.Type ZDataType }

emptyEnv = HsEnv    { envVars = mempty,
                      envTypes = mempty }
                      
type HsZeno = StateT HsEnv ZenoState
type HsExpr = Hs.Expr Hs.Var
type HsBinding = (Hs.Var, HsExpr)

instance Eq Hs.Type where
  (==) = Hs.tcEqType
                     
instance Ord Hs.Type where
  compare = Hs.tcCmpType
  
parseHaskellFile :: FilePath -> IO (ZenoState ())
parseHaskellFile filepath = do
  core_mod <- Hs.defaultErrorHandler Hs.defaultDynFlags $ 
    Hs.runGhc (Just Paths.libdir) $ do
      dflags <- Hs.getSessionDynFlags
      Hs.setSessionDynFlags $ dflags 
      {-{ Hs.hscTarget = Hs.HscInterpreted, 
          Hs.ghcLink = Hs.LinkInMemory } -}
      Hs.compileToCoreSimplified filepath
      
  --putStrLn (output core_mod)
      
  let types = Hs.cm_types core_mod
      binds = Hs.flattenBinds (Hs.cm_binds core_mod)
  
  return $ do
    flip evalStateT emptyEnv $ do
      convertDataCons (Hs.typeEnvDataCons types)
      func_defs <- mapM convertBinding binds
      lift $ mapM_ addFunctionDef func_defs
      
  where
    isLemmaBinding :: HsBinding -> Bool
    isLemmaBinding (var, _)
      | output (resultHsType (Hs.varType var)) == "Z.Lemma" = True
    isLemmaBinding _ = False

addEnvVar :: Hs.Var -> ZVar -> HsZeno ()
addEnvVar hs_var z_var = modify $ \env -> 
  env { envVars = Map.insert hs_var z_var (envVars env) }
  
resultHsType :: Hs.Type -> Hs.Type
resultHsType t = 
  case Hs.splitFunTy_maybe t of
    Nothing -> t
    Just (_, t') -> resultHsType t'

stripModuleName :: String -> String
stripModuleName name = 
  if '.' `elem` name 
    then tail $ dropWhile (/= '.') name
    else name
    
createEnvVar :: Hs.Var -> ZVarClass -> HsZeno ZVar
createEnvVar hsvar var_cls = do
  new_id <- lift newIdS
  var_type <- convertCoreType (Hs.varType hsvar)
  let new_var = ZVar new_id var_type var_cls
  addEnvVar hsvar new_var
  lift $ addGlobalVar new_var (outputName hsvar) 
  return new_var
  
addEnvType :: Hs.Type -> ZDataType -> HsZeno ()
addEnvType hs_type dtype = modify $ \env -> 
  env { envTypes = Map.insert hs_type dtype (envTypes env) }
  
output :: Outputable a => a -> String
output = flip ppr defaultUserStyle >>> show

outputName :: Outputable a => a -> String
outputName = stripModuleName . output

lookupTypeId :: Hs.Type -> HsZeno ZDataType
lookupTypeId coretype = do
  types <- envTypes <$> get
  case Map.lookup coretype types of
    Just dtype -> return dtype
    
lookupVar :: Hs.Var -> HsZeno ZVar 
lookupVar hsvar = do
  vars <- envVars <$> get
  case Map.lookup hsvar vars of
    Just zvar -> return zvar

convertCoreType :: Hs.Type -> HsZeno ZType
convertCoreType coretype = do
  case Hs.splitForAllTy_maybe coretype of
    Just (all, inner) ->
      convertCoreType inner
    Nothing ->
      case Hs.splitFunTy_maybe coretype of
        Just (arg, res) -> do
          z_arg <- convertCoreType arg
          z_res <- convertCoreType res
          return $ Func z_arg z_res
        Nothing ->
          Type <$> (lookupTypeId coretype)

convertBinding :: HsBinding -> HsZeno ZFunctionDef
convertBinding (hsvar, expr) = do
  fun_var <- createEnvVar hsvar ZFunction
  (args, stripped_expr) <- stripArgs expr
  fun_expr <- convertExpr stripped_expr
  return $ FunctionDef fun_var args fun_expr
  where
    stripArgs :: HsExpr -> HsZeno ([ZVar], HsExpr)
    stripArgs (Hs.Lam hs_var expr) = do
      z_var <- createEnvVar hs_var ZVariable
      addEnvVar hs_var z_var
      (vars, expr') <- stripArgs expr
      return (z_var : vars, expr')
    stripArgs expr = return ([], expr)

convertTermExpr :: HsExpr -> HsZeno ZTerm
convertTermExpr (Hs.Var hsvar) =
  Var <$> (lookupVar hsvar)
convertTermExpr (Hs.App fun arg) = do
  zfun <- convertTermExpr fun
  zarg <- convertTermExpr arg
  return $ App zfun zarg
  
mapHsExpr :: (HsExpr -> HsExpr) -> HsExpr -> HsExpr
mapHsExpr f (Hs.App left right) = f (Hs.App (f left) (f right))
mapHsExpr f (Hs.Lam b expr) = f (Hs.Lam b (f expr)) 
mapHsExpr f (Hs.Let binds expr) =
  let mapBind (Hs.NonRec b expr) = Hs.NonRec b (f expr)
  in f (Hs.Let (mapBind binds) (f expr))
mapHsExpr f (Hs.Case expr b t alts) =
  let mapAlt (con, bs, expr) = (con, bs, f expr)
  in f (Hs.Case (f expr) b t (map mapAlt alts))
mapHsExpr f (Hs.Cast expr coer) = f (Hs.Cast (f expr) coer)
mapHsExpr f (Hs.Note note expr) = f (Hs.Note note (f expr))
mapHsExpr f expr = f expr
    
convertExpr :: HsExpr -> HsZeno ZExpression
convertExpr (Hs.Case of_expr hsvar _ alts) = do
  z_var <- createEnvVar hsvar ZVariable
  whens' <- mapM convertAlts alts
  of_term <- convertTermExpr of_expr
  let whens = map (replaceTerms (Var z_var) of_term) whens' 
  return $ CaseOf of_term whens
convertExpr expr =
  Value <$> convertTermExpr expr
  
convertAlts :: Hs.Alt Hs.Var -> HsZeno ZWhen 
convertAlts (Hs.DataAlt con, binds, expr) = do
  con_var <- lookupVar (Hs.dataConWorkId con)
  bind_vars <- mapM (flip createEnvVar ZVariable) binds
  z_expr <- convertExpr expr
  let term = unflattenTerm $ map Var (con_var : bind_vars)
  return $ When term z_expr 
  
convertDataCons :: [Hs.DataCon] -> HsZeno ()
convertDataCons cons = do
  let con_types = nub $ map dataConReturnType cons
  
      typed_cons :: [(Hs.Type, [Hs.DataCon])]
      typed_cons = map (\t -> (t, consOfType t)) con_types
      
  rec types <- mapM (createDataType types) typed_cons
  return ()
  where
    createDataType :: [(Hs.Type, ZDataType)] 
      -> (Hs.Type, [Hs.DataCon]) -> HsZeno (Hs.Type, ZDataType)
    createDataType knot (typ, cons) = do
      type_id <- lift newIdS
      rec let datatype = ZDataType type_id cons' 
          cons' <- mapM (createDataConVar datatype) cons
      addEnvType typ datatype
      lift $ do addDataType datatype
                christen type_id (outputName typ)
      return (typ, datatype)
      where
        -- No function types allowed yet
        createDataConVar :: ZDataType -> Hs.DataCon -> HsZeno ZVar
        createDataConVar ret_type con = do
          var_id <- lift newIdS
          
          let arg_types :: [ZType]
              arg_types = map 
                (\t -> Type $ fromJust $ lookup t knot) 
                (Hs.dataConOrigArgTys con)
              var_type = unflattenType (arg_types ++ [Type ret_type])  
              new_var = ZVar var_id var_type ZConstructor
              
          lift $ addGlobalVar new_var 
            (outputName $ Hs.dataConName con)
          addEnvVar (Hs.dataConWorkId con) new_var
          return new_var
          
    consOfType :: Hs.Type -> [Hs.DataCon]
    consOfType t = filter ((== t) . dataConReturnType) cons
  
    dataConReturnType :: Hs.DataCon -> Hs.Type
    dataConReturnType con = 
      let (_, _, _, t) = Hs.dataConSig con in t
      
      
      
      
